#include "TcpServer.h"
#include "Acceptor.h"
#include "EventLoop.h"
#include "Logger.h"

namespace GiSo {

class TcpConnection;

TcpServer::TcpServer(const InetAddress& addr, int threadNum)
    : m_EventLoop(threadNum), m_Address(addr), m_Acceptor(new Acceptor(&m_EventLoop, addr)) {
    m_EachThreadConnectionCount.resize(threadNum, 0);

    m_Acceptor->SetNewConnectionCallback([this](std::unique_ptr<Channel> channel, const InetAddress& peerAddr) {
        int i = GetMinConnectionThread();

        auto conn = std::make_shared<TcpConnection>(std::move(channel), m_EventLoop.GetTaskQueue(i));

        if (m_NewConnectionCallback && m_NewConnectionCallback(conn, peerAddr)) {
            AddConnection(conn->GetFd(), i, conn);
            GS_TRACE("Add Connection to the TcpServer - fd={0}", conn->GetFd());

            // 这里之所以要把remove操作放到主线程中，是因为poller返回的是channel的原始指针，而不是shared_ptr
            // 这就导致了如果不控制remove的顺序，有可能原始指针指向的内容已经被析构了
            // 所以这里要把remove操作放到主线程中，通过容器中存放的shared_ptr来保证其生命周期
            conn->SetRemoveFromContainerCallback([this](std::shared_ptr<TcpConnection> conn) {
                this->m_EventLoop.GetMainQueue()->PushBack([this,conn]() {
                    RemoveConnection(conn->GetFd());
                });
            });
        } else {
            GS_TRACE("Reject the Connection - fd={0}", conn->GetFd());
            conn->Close();
        }
    });
}

TcpServer::~TcpServer() {
    Close();
}

void TcpServer::Start() {
    m_Acceptor->Listen();
    m_EventLoop.Run();
}

void TcpServer::Stop() {
    m_Running = false;
    m_EventLoop.Stop();
}

void TcpServer::Close() {
    if (m_Running) {
        Stop();
    }
    for (auto& [fd, conn] : m_Connections) {
        conn.second->Close();
    }
    m_Connections.clear();
}

void TcpServer::SetNewConnectionCallback(std::function<bool(std::shared_ptr<TcpConnection>, const InetAddress&)> cb) {
    m_NewConnectionCallback = std::move(cb);
}

void TcpServer::SetCloseConnectionCallback(std::function<void(std::shared_ptr<TcpConnection>)> cb) {
    m_CloseConnectionCallback = std::move(cb);
}

int TcpServer::GetMinConnectionThread() {
    std::lock_guard<std::mutex> lock(m_Mutex);
    int min = m_EachThreadConnectionCount[0];
    int minIndex = 0;
    for (int i = 1; i < m_EachThreadConnectionCount.size(); i++) {
        if (m_EachThreadConnectionCount[i] < min) {
            min = m_EachThreadConnectionCount[i];
            minIndex = i;
        }
    }
    return minIndex;
}

void TcpServer::AddConnection(int id, int threadId, std::shared_ptr<TcpConnection> conn) {
    std::lock_guard<std::mutex> lock(m_Mutex);
    m_Connections[id] = {threadId, conn};
    m_EachThreadConnectionCount[threadId]++;
}

void TcpServer::RemoveConnection(int id) {
    std::lock_guard<std::mutex> lock(m_Mutex);
    auto it = m_Connections.find(id);
    if (it == m_Connections.end()) {
        return;
    }
    auto threadId = it->second.first;
    int n = m_Connections.erase(id);
    if (n > 0) {
        m_EachThreadConnectionCount[threadId]--;
    }
}

}